#!/usr/bin/env python3
import re
import os

def shard_markdown(input_file, output_dir):
    """Shard a markdown document by level 2 sections."""
    
    # Read the entire document
    with open(input_file, 'r', encoding='utf-8') as f:
        content = f.read()
    
    # Create output directory if it doesn't exist
    os.makedirs(output_dir, exist_ok=True)
    
    # Find the main title and content before first level 2 section
    lines = content.split('\n')
    main_title = ""
    intro_content = []
    current_line = 0
    
    # Find the main title (first level 1 heading)
    for i, line in enumerate(lines):
        if line.startswith('# '):
            main_title = line[2:].strip()
            current_line = i + 1
            break
    
    # Collect content before first level 2 section
    in_code_block = False
    for i in range(current_line, len(lines)):
        line = lines[i]
        
        # Track code blocks
        if line.strip().startswith('```'):
            in_code_block = not in_code_block
        
        # Stop at first level 2 heading (but not if in code block)
        if line.startswith('## ') and not in_code_block:
            current_line = i
            break
        
        intro_content.append(line)
    
    # Parse sections
    sections = []
    current_section = None
    in_code_block = False
    
    for i in range(current_line, len(lines)):
        line = lines[i]
        
        # Track code blocks
        if line.strip().startswith('```'):
            in_code_block = not in_code_block
        
        # Check for level 2 heading (but not if in code block)
        if line.startswith('## ') and not in_code_block:
            # Save previous section if exists
            if current_section:
                sections.append(current_section)
            
            # Start new section
            current_section = {
                'title': line[3:].strip(),
                'content': [line]
            }
        elif current_section:
            current_section['content'].append(line)
    
    # Don't forget the last section
    if current_section:
        sections.append(current_section)
    
    # Process each section
    section_files = []
    for section in sections:
        # Generate filename
        filename = section['title'].lower()
        # Remove section numbers if present (e.g., "1. " or "10. ")
        filename = re.sub(r'^\d+\.\s+', '', filename)
        # Convert to lowercase-dash-case
        filename = re.sub(r'[^a-z0-9]+', '-', filename)
        filename = filename.strip('-') + '.md'
        
        # Adjust heading levels in content
        adjusted_content = []
        for line in section['content']:
            # Skip if in code block
            if line.strip().startswith('```'):
                in_code_block = not in_code_block
                adjusted_content.append(line)
            elif not in_code_block:
                # Adjust heading levels
                if line.startswith('## '):
                    adjusted_content.append('# ' + line[3:])
                elif line.startswith('### '):
                    adjusted_content.append('## ' + line[4:])
                elif line.startswith('#### '):
                    adjusted_content.append('### ' + line[5:])
                elif line.startswith('##### '):
                    adjusted_content.append('#### ' + line[6:])
                elif line.startswith('###### '):
                    adjusted_content.append('##### ' + line[7:])
                else:
                    adjusted_content.append(line)
            else:
                adjusted_content.append(line)
        
        # Write to file
        filepath = os.path.join(output_dir, filename)
        with open(filepath, 'w', encoding='utf-8') as f:
            f.write('\n'.join(adjusted_content))
        
        section_files.append({
            'filename': filename,
            'title': section['title']
        })
    
    # Create index.md
    index_content = []
    if main_title:
        index_content.append(f"# {main_title}")
        index_content.append("")
    
    # Add intro content if exists
    if intro_content and any(line.strip() for line in intro_content):
        index_content.extend(intro_content)
        index_content.append("")
    
    # Add sections list
    index_content.append("## Sections")
    index_content.append("")
    for section_file in section_files:
        index_content.append(f"- [{section_file['title']}](./{section_file['filename']})")
    
    # Write index.md
    index_path = os.path.join(output_dir, 'index.md')
    with open(index_path, 'w', encoding='utf-8') as f:
        f.write('\n'.join(index_content))
    
    return len(section_files)

# Run the sharding
if __name__ == "__main__":
    input_file = "/mnt/d/aiProject/paper/aipaper/docs/architecture.md"
    output_dir = "/mnt/d/aiProject/paper/aipaper/docs/architecture"
    
    num_files = shard_markdown(input_file, output_dir)
    print(f"Document sharded successfully!")
    print(f"- Source: {input_file}")
    print(f"- Destination: {output_dir}/")
    print(f"- Files created: {num_files + 1} (including index.md)")